from enhanced_classifier import EnhancedClassifier

def main():
    # 准备训练数据
    training_documents = [
        "这是一篇关于人工智能的文章，讨论了机器学习和深度学习的应用",
        "今天的天气很好，阳光明媚，适合出去散步",
        "最新的手机采用了先进的人工智能技术，能够实现智能语音助手功能",
        "周末去公园野餐，遇到了很多朋友，玩得很开心",
        "深度学习技术在图像识别领域取得了突破性进展",
        "春天来了，花园里的花都开了，真是美丽",
        "研究人员开发出新的自然语言处理模型",
        "昨天参加了一场精彩的音乐会，感觉很放松"
    ]
    
    training_labels = ['技术', '生活', '技术', '生活', '技术', '生活', '技术', '生活']
    
    # 准备测试数据
    test_documents = [
        "新开发的人工智能系统可以自动进行图像识别",
        "昨天和家人一起去海边度假，拍了很多照片",
        "量子计算机在机器学习领域展现出巨大潜力",
        "周末在家做了一顿美味的晚餐"
    ]
    test_labels = ['技术', '生活', '技术', '生活']
    
    # 创建增强版分类器实例
    classifier = EnhancedClassifier(vector_size=100, window=5, min_count=1)
    
    # 训练分类器
    print("正在训练分类器...")
    print("1. 训练Word2Vec模型")
    print("2. 计算TF-IDF权重")
    print("3. 训练SVM分类器")
    classifier.train(training_documents, training_labels)
    
    # 评估分类器性能
    print("\n分类器评估报告:")
    print(classifier.evaluate(test_documents, test_labels))
    
    # 测试新文档分类
    new_documents = [
        "这款新手机采用了5nm制程的处理器，性能提升显著",
        "小区里新开了一家面包店，每天都飘着香味"
    ]
    
    print("\n新文档分类结果:")
    for doc in new_documents:
        predicted_category = classifier.predict(doc)
        print(f"文档: {doc}")
        print(f"预测类别: {predicted_category}\n")

if __name__ == '__main__':
    main()